{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "B8a_URGiowPn"
      },
      "source": [
        "## Overview\n",
        "This colab demonstrates the steps to run a family of DeepLab models built by the DeepLab2 library to perform dense pixel labeling tasks. The models used in this colab perform panoptic segmentation, where the predicted value encodes both semantic class and instance label for every pixel (including both ‘thing’ and ‘stuff’ pixels).\n",
        "\n",
        "### About DeepLab2\n",
        "DeepLab2 is a TensorFlow library for deep labeling, aiming to facilitate future research on dense pixel labeling tasks by providing state-of-the-art and easy-to-use TensorFlow models. Code is made publicly available at https://github.com/google-research/deeplab2\n",
        "\n",
        "The checkpoints used in this demo are pretrained on the COCO dataset."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IGVFjkE2o0K8"
      },
      "source": [
        "### Import and helper methods"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "dQNiIp-LoV6f"
      },
      "outputs": [],
      "source": [
        "import collections\n",
        "import os\n",
        "import tempfile\n",
        "import copy\n",
        "\n",
        "from matplotlib import gridspec\n",
        "from matplotlib import pyplot as plt\n",
        "import numpy as np\n",
        "from PIL import Image\n",
        "import urllib\n",
        "\n",
        "import tensorflow as tf\n",
        "\n",
        "from google.colab import files"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "0CF8vyASXwEr"
      },
      "outputs": [],
      "source": [
        "COCO_META = [\n",
        "    {\n",
        "        'color': [220, 20, 60],\n",
        "        'isthing': 1,\n",
        "        'id': 1,\n",
        "        'name': 'person'\n",
        "    },\n",
        "    {\n",
        "        'color': [119, 11, 32],\n",
        "        'isthing': 1,\n",
        "        'id': 2,\n",
        "        'name': 'bicycle'\n",
        "    },\n",
        "    {\n",
        "        'color': [0, 0, 142],\n",
        "        'isthing': 1,\n",
        "        'id': 3,\n",
        "        'name': 'car'\n",
        "    },\n",
        "    {\n",
        "        'color': [0, 0, 230],\n",
        "        'isthing': 1,\n",
        "        'id': 4,\n",
        "        'name': 'motorcycle'\n",
        "    },\n",
        "    {\n",
        "        'color': [106, 0, 228],\n",
        "        'isthing': 1,\n",
        "        'id': 5,\n",
        "        'name': 'airplane'\n",
        "    },\n",
        "    {\n",
        "        'color': [0, 60, 100],\n",
        "        'isthing': 1,\n",
        "        'id': 6,\n",
        "        'name': 'bus'\n",
        "    },\n",
        "    {\n",
        "        'color': [0, 80, 100],\n",
        "        'isthing': 1,\n",
        "        'id': 7,\n",
        "        'name': 'train'\n",
        "    },\n",
        "    {\n",
        "        'color': [0, 0, 70],\n",
        "        'isthing': 1,\n",
        "        'id': 8,\n",
        "        'name': 'truck'\n",
        "    },\n",
        "    {\n",
        "        'color': [0, 0, 192],\n",
        "        'isthing': 1,\n",
        "        'id': 9,\n",
        "        'name': 'boat'\n",
        "    },\n",
        "    {\n",
        "        'color': [250, 170, 30],\n",
        "        'isthing': 1,\n",
        "        'id': 10,\n",
        "        'name': 'traffic light'\n",
        "    },\n",
        "    {\n",
        "        'color': [100, 170, 30],\n",
        "        'isthing': 1,\n",
        "        'id': 11,\n",
        "        'name': 'fire hydrant'\n",
        "    },\n",
        "    {\n",
        "        'color': [220, 220, 0],\n",
        "        'isthing': 1,\n",
        "        'id': 13,\n",
        "        'name': 'stop sign'\n",
        "    },\n",
        "    {\n",
        "        'color': [175, 116, 175],\n",
        "        'isthing': 1,\n",
        "        'id': 14,\n",
        "        'name': 'parking meter'\n",
        "    },\n",
        "    {\n",
        "        'color': [250, 0, 30],\n",
        "        'isthing': 1,\n",
        "        'id': 15,\n",
        "        'name': 'bench'\n",
        "    },\n",
        "    {\n",
        "        'color': [165, 42, 42],\n",
        "        'isthing': 1,\n",
        "        'id': 16,\n",
        "        'name': 'bird'\n",
        "    },\n",
        "    {\n",
        "        'color': [255, 77, 255],\n",
        "        'isthing': 1,\n",
        "        'id': 17,\n",
        "        'name': 'cat'\n",
        "    },\n",
        "    {\n",
        "        'color': [0, 226, 252],\n",
        "        'isthing': 1,\n",
        "        'id': 18,\n",
        "        'name': 'dog'\n",
        "    },\n",
        "    {\n",
        "        'color': [182, 182, 255],\n",
        "        'isthing': 1,\n",
        "        'id': 19,\n",
        "        'name': 'horse'\n",
        "    },\n",
        "    {\n",
        "        'color': [0, 82, 0],\n",
        "        'isthing': 1,\n",
        "        'id': 20,\n",
        "        'name': 'sheep'\n",
        "    },\n",
        "    {\n",
        "        'color': [120, 166, 157],\n",
        "        'isthing': 1,\n",
        "        'id': 21,\n",
        "        'name': 'cow'\n",
        "    },\n",
        "    {\n",
        "        'color': [110, 76, 0],\n",
        "        'isthing': 1,\n",
        "        'id': 22,\n",
        "        'name': 'elephant'\n",
        "    },\n",
        "    {\n",
        "        'color': [174, 57, 255],\n",
        "        'isthing': 1,\n",
        "        'id': 23,\n",
        "        'name': 'bear'\n",
        "    },\n",
        "    {\n",
        "        'color': [199, 100, 0],\n",
        "        'isthing': 1,\n",
        "        'id': 24,\n",
        "        'name': 'zebra'\n",
        "    },\n",
        "    {\n",
        "        'color': [72, 0, 118],\n",
        "        'isthing': 1,\n",
        "        'id': 25,\n",
        "        'name': 'giraffe'\n",
        "    },\n",
        "    {\n",
        "        'color': [255, 179, 240],\n",
        "        'isthing': 1,\n",
        "        'id': 27,\n",
        "        'name': 'backpack'\n",
        "    },\n",
        "    {\n",
        "        'color': [0, 125, 92],\n",
        "        'isthing': 1,\n",
        "        'id': 28,\n",
        "        'name': 'umbrella'\n",
        "    },\n",
        "    {\n",
        "        'color': [209, 0, 151],\n",
        "        'isthing': 1,\n",
        "        'id': 31,\n",
        "        'name': 'handbag'\n",
        "    },\n",
        "    {\n",
        "        'color': [188, 208, 182],\n",
        "        'isthing': 1,\n",
        "        'id': 32,\n",
        "        'name': 'tie'\n",
        "    },\n",
        "    {\n",
        "        'color': [0, 220, 176],\n",
        "        'isthing': 1,\n",
        "        'id': 33,\n",
        "        'name': 'suitcase'\n",
        "    },\n",
        "    {\n",
        "        'color': [255, 99, 164],\n",
        "        'isthing': 1,\n",
        "        'id': 34,\n",
        "        'name': 'frisbee'\n",
        "    },\n",
        "    {\n",
        "        'color': [92, 0, 73],\n",
        "        'isthing': 1,\n",
        "        'id': 35,\n",
        "        'name': 'skis'\n",
        "    },\n",
        "    {\n",
        "        'color': [133, 129, 255],\n",
        "        'isthing': 1,\n",
        "        'id': 36,\n",
        "        'name': 'snowboard'\n",
        "    },\n",
        "    {\n",
        "        'color': [78, 180, 255],\n",
        "        'isthing': 1,\n",
        "        'id': 37,\n",
        "        'name': 'sports ball'\n",
        "    },\n",
        "    {\n",
        "        'color': [0, 228, 0],\n",
        "        'isthing': 1,\n",
        "        'id': 38,\n",
        "        'name': 'kite'\n",
        "    },\n",
        "    {\n",
        "        'color': [174, 255, 243],\n",
        "        'isthing': 1,\n",
        "        'id': 39,\n",
        "        'name': 'baseball bat'\n",
        "    },\n",
        "    {\n",
        "        'color': [45, 89, 255],\n",
        "        'isthing': 1,\n",
        "        'id': 40,\n",
        "        'name': 'baseball glove'\n",
        "    },\n",
        "    {\n",
        "        'color': [134, 134, 103],\n",
        "        'isthing': 1,\n",
        "        'id': 41,\n",
        "        'name': 'skateboard'\n",
        "    },\n",
        "    {\n",
        "        'color': [145, 148, 174],\n",
        "        'isthing': 1,\n",
        "        'id': 42,\n",
        "        'name': 'surfboard'\n",
        "    },\n",
        "    {\n",
        "        'color': [255, 208, 186],\n",
        "        'isthing': 1,\n",
        "        'id': 43,\n",
        "        'name': 'tennis racket'\n",
        "    },\n",
        "    {\n",
        "        'color': [197, 226, 255],\n",
        "        'isthing': 1,\n",
        "        'id': 44,\n",
        "        'name': 'bottle'\n",
        "    },\n",
        "    {\n",
        "        'color': [171, 134, 1],\n",
        "        'isthing': 1,\n",
        "        'id': 46,\n",
        "        'name': 'wine glass'\n",
        "    },\n",
        "    {\n",
        "        'color': [109, 63, 54],\n",
        "        'isthing': 1,\n",
        "        'id': 47,\n",
        "        'name': 'cup'\n",
        "    },\n",
        "    {\n",
        "        'color': [207, 138, 255],\n",
        "        'isthing': 1,\n",
        "        'id': 48,\n",
        "        'name': 'fork'\n",
        "    },\n",
        "    {\n",
        "        'color': [151, 0, 95],\n",
        "        'isthing': 1,\n",
        "        'id': 49,\n",
        "        'name': 'knife'\n",
        "    },\n",
        "    {\n",
        "        'color': [9, 80, 61],\n",
        "        'isthing': 1,\n",
        "        'id': 50,\n",
        "        'name': 'spoon'\n",
        "    },\n",
        "    {\n",
        "        'color': [84, 105, 51],\n",
        "        'isthing': 1,\n",
        "        'id': 51,\n",
        "        'name': 'bowl'\n",
        "    },\n",
        "    {\n",
        "        'color': [74, 65, 105],\n",
        "        'isthing': 1,\n",
        "        'id': 52,\n",
        "        'name': 'banana'\n",
        "    },\n",
        "    {\n",
        "        'color': [166, 196, 102],\n",
        "        'isthing': 1,\n",
        "        'id': 53,\n",
        "        'name': 'apple'\n",
        "    },\n",
        "    {\n",
        "        'color': [208, 195, 210],\n",
        "        'isthing': 1,\n",
        "        'id': 54,\n",
        "        'name': 'sandwich'\n",
        "    },\n",
        "    {\n",
        "        'color': [255, 109, 65],\n",
        "        'isthing': 1,\n",
        "        'id': 55,\n",
        "        'name': 'orange'\n",
        "    },\n",
        "    {\n",
        "        'color': [0, 143, 149],\n",
        "        'isthing': 1,\n",
        "        'id': 56,\n",
        "        'name': 'broccoli'\n",
        "    },\n",
        "    {\n",
        "        'color': [179, 0, 194],\n",
        "        'isthing': 1,\n",
        "        'id': 57,\n",
        "        'name': 'carrot'\n",
        "    },\n",
        "    {\n",
        "        'color': [209, 99, 106],\n",
        "        'isthing': 1,\n",
        "        'id': 58,\n",
        "        'name': 'hot dog'\n",
        "    },\n",
        "    {\n",
        "        'color': [5, 121, 0],\n",
        "        'isthing': 1,\n",
        "        'id': 59,\n",
        "        'name': 'pizza'\n",
        "    },\n",
        "    {\n",
        "        'color': [227, 255, 205],\n",
        "        'isthing': 1,\n",
        "        'id': 60,\n",
        "        'name': 'donut'\n",
        "    },\n",
        "    {\n",
        "        'color': [147, 186, 208],\n",
        "        'isthing': 1,\n",
        "        'id': 61,\n",
        "        'name': 'cake'\n",
        "    },\n",
        "    {\n",
        "        'color': [153, 69, 1],\n",
        "        'isthing': 1,\n",
        "        'id': 62,\n",
        "        'name': 'chair'\n",
        "    },\n",
        "    {\n",
        "        'color': [3, 95, 161],\n",
        "        'isthing': 1,\n",
        "        'id': 63,\n",
        "        'name': 'couch'\n",
        "    },\n",
        "    {\n",
        "        'color': [163, 255, 0],\n",
        "        'isthing': 1,\n",
        "        'id': 64,\n",
        "        'name': 'potted plant'\n",
        "    },\n",
        "    {\n",
        "        'color': [119, 0, 170],\n",
        "        'isthing': 1,\n",
        "        'id': 65,\n",
        "        'name': 'bed'\n",
        "    },\n",
        "    {\n",
        "        'color': [0, 182, 199],\n",
        "        'isthing': 1,\n",
        "        'id': 67,\n",
        "        'name': 'dining table'\n",
        "    },\n",
        "    {\n",
        "        'color': [0, 165, 120],\n",
        "        'isthing': 1,\n",
        "        'id': 70,\n",
        "        'name': 'toilet'\n",
        "    },\n",
        "    {\n",
        "        'color': [183, 130, 88],\n",
        "        'isthing': 1,\n",
        "        'id': 72,\n",
        "        'name': 'tv'\n",
        "    },\n",
        "    {\n",
        "        'color': [95, 32, 0],\n",
        "        'isthing': 1,\n",
        "        'id': 73,\n",
        "        'name': 'laptop'\n",
        "    },\n",
        "    {\n",
        "        'color': [130, 114, 135],\n",
        "        'isthing': 1,\n",
        "        'id': 74,\n",
        "        'name': 'mouse'\n",
        "    },\n",
        "    {\n",
        "        'color': [110, 129, 133],\n",
        "        'isthing': 1,\n",
        "        'id': 75,\n",
        "        'name': 'remote'\n",
        "    },\n",
        "    {\n",
        "        'color': [166, 74, 118],\n",
        "        'isthing': 1,\n",
        "        'id': 76,\n",
        "        'name': 'keyboard'\n",
        "    },\n",
        "    {\n",
        "        'color': [219, 142, 185],\n",
        "        'isthing': 1,\n",
        "        'id': 77,\n",
        "        'name': 'cell phone'\n",
        "    },\n",
        "    {\n",
        "        'color': [79, 210, 114],\n",
        "        'isthing': 1,\n",
        "        'id': 78,\n",
        "        'name': 'microwave'\n",
        "    },\n",
        "    {\n",
        "        'color': [178, 90, 62],\n",
        "        'isthing': 1,\n",
        "        'id': 79,\n",
        "        'name': 'oven'\n",
        "    },\n",
        "    {\n",
        "        'color': [65, 70, 15],\n",
        "        'isthing': 1,\n",
        "        'id': 80,\n",
        "        'name': 'toaster'\n",
        "    },\n",
        "    {\n",
        "        'color': [127, 167, 115],\n",
        "        'isthing': 1,\n",
        "        'id': 81,\n",
        "        'name': 'sink'\n",
        "    },\n",
        "    {\n",
        "        'color': [59, 105, 106],\n",
        "        'isthing': 1,\n",
        "        'id': 82,\n",
        "        'name': 'refrigerator'\n",
        "    },\n",
        "    {\n",
        "        'color': [142, 108, 45],\n",
        "        'isthing': 1,\n",
        "        'id': 84,\n",
        "        'name': 'book'\n",
        "    },\n",
        "    {\n",
        "        'color': [196, 172, 0],\n",
        "        'isthing': 1,\n",
        "        'id': 85,\n",
        "        'name': 'clock'\n",
        "    },\n",
        "    {\n",
        "        'color': [95, 54, 80],\n",
        "        'isthing': 1,\n",
        "        'id': 86,\n",
        "        'name': 'vase'\n",
        "    },\n",
        "    {\n",
        "        'color': [128, 76, 255],\n",
        "        'isthing': 1,\n",
        "        'id': 87,\n",
        "        'name': 'scissors'\n",
        "    },\n",
        "    {\n",
        "        'color': [201, 57, 1],\n",
        "        'isthing': 1,\n",
        "        'id': 88,\n",
        "        'name': 'teddy bear'\n",
        "    },\n",
        "    {\n",
        "        'color': [246, 0, 122],\n",
        "        'isthing': 1,\n",
        "        'id': 89,\n",
        "        'name': 'hair drier'\n",
        "    },\n",
        "    {\n",
        "        'color': [191, 162, 208],\n",
        "        'isthing': 1,\n",
        "        'id': 90,\n",
        "        'name': 'toothbrush'\n",
        "    },\n",
        "    {\n",
        "        'color': [255, 255, 128],\n",
        "        'isthing': 0,\n",
        "        'id': 92,\n",
        "        'name': 'banner'\n",
        "    },\n",
        "    {\n",
        "        'color': [147, 211, 203],\n",
        "        'isthing': 0,\n",
        "        'id': 93,\n",
        "        'name': 'blanket'\n",
        "    },\n",
        "    {\n",
        "        'color': [150, 100, 100],\n",
        "        'isthing': 0,\n",
        "        'id': 95,\n",
        "        'name': 'bridge'\n",
        "    },\n",
        "    {\n",
        "        'color': [168, 171, 172],\n",
        "        'isthing': 0,\n",
        "        'id': 100,\n",
        "        'name': 'cardboard'\n",
        "    },\n",
        "    {\n",
        "        'color': [146, 112, 198],\n",
        "        'isthing': 0,\n",
        "        'id': 107,\n",
        "        'name': 'counter'\n",
        "    },\n",
        "    {\n",
        "        'color': [210, 170, 100],\n",
        "        'isthing': 0,\n",
        "        'id': 109,\n",
        "        'name': 'curtain'\n",
        "    },\n",
        "    {\n",
        "        'color': [92, 136, 89],\n",
        "        'isthing': 0,\n",
        "        'id': 112,\n",
        "        'name': 'door-stuff'\n",
        "    },\n",
        "    {\n",
        "        'color': [218, 88, 184],\n",
        "        'isthing': 0,\n",
        "        'id': 118,\n",
        "        'name': 'floor-wood'\n",
        "    },\n",
        "    {\n",
        "        'color': [241, 129, 0],\n",
        "        'isthing': 0,\n",
        "        'id': 119,\n",
        "        'name': 'flower'\n",
        "    },\n",
        "    {\n",
        "        'color': [217, 17, 255],\n",
        "        'isthing': 0,\n",
        "        'id': 122,\n",
        "        'name': 'fruit'\n",
        "    },\n",
        "    {\n",
        "        'color': [124, 74, 181],\n",
        "        'isthing': 0,\n",
        "        'id': 125,\n",
        "        'name': 'gravel'\n",
        "    },\n",
        "    {\n",
        "        'color': [70, 70, 70],\n",
        "        'isthing': 0,\n",
        "        'id': 128,\n",
        "        'name': 'house'\n",
        "    },\n",
        "    {\n",
        "        'color': [255, 228, 255],\n",
        "        'isthing': 0,\n",
        "        'id': 130,\n",
        "        'name': 'light'\n",
        "    },\n",
        "    {\n",
        "        'color': [154, 208, 0],\n",
        "        'isthing': 0,\n",
        "        'id': 133,\n",
        "        'name': 'mirror-stuff'\n",
        "    },\n",
        "    {\n",
        "        'color': [193, 0, 92],\n",
        "        'isthing': 0,\n",
        "        'id': 138,\n",
        "        'name': 'net'\n",
        "    },\n",
        "    {\n",
        "        'color': [76, 91, 113],\n",
        "        'isthing': 0,\n",
        "        'id': 141,\n",
        "        'name': 'pillow'\n",
        "    },\n",
        "    {\n",
        "        'color': [255, 180, 195],\n",
        "        'isthing': 0,\n",
        "        'id': 144,\n",
        "        'name': 'platform'\n",
        "    },\n",
        "    {\n",
        "        'color': [106, 154, 176],\n",
        "        'isthing': 0,\n",
        "        'id': 145,\n",
        "        'name': 'playingfield'\n",
        "    },\n",
        "    {\n",
        "        'color': [230, 150, 140],\n",
        "        'isthing': 0,\n",
        "        'id': 147,\n",
        "        'name': 'railroad'\n",
        "    },\n",
        "    {\n",
        "        'color': [60, 143, 255],\n",
        "        'isthing': 0,\n",
        "        'id': 148,\n",
        "        'name': 'river'\n",
        "    },\n",
        "    {\n",
        "        'color': [128, 64, 128],\n",
        "        'isthing': 0,\n",
        "        'id': 149,\n",
        "        'name': 'road'\n",
        "    },\n",
        "    {\n",
        "        'color': [92, 82, 55],\n",
        "        'isthing': 0,\n",
        "        'id': 151,\n",
        "        'name': 'roof'\n",
        "    },\n",
        "    {\n",
        "        'color': [254, 212, 124],\n",
        "        'isthing': 0,\n",
        "        'id': 154,\n",
        "        'name': 'sand'\n",
        "    },\n",
        "    {\n",
        "        'color': [73, 77, 174],\n",
        "        'isthing': 0,\n",
        "        'id': 155,\n",
        "        'name': 'sea'\n",
        "    },\n",
        "    {\n",
        "        'color': [255, 160, 98],\n",
        "        'isthing': 0,\n",
        "        'id': 156,\n",
        "        'name': 'shelf'\n",
        "    },\n",
        "    {\n",
        "        'color': [255, 255, 255],\n",
        "        'isthing': 0,\n",
        "        'id': 159,\n",
        "        'name': 'snow'\n",
        "    },\n",
        "    {\n",
        "        'color': [104, 84, 109],\n",
        "        'isthing': 0,\n",
        "        'id': 161,\n",
        "        'name': 'stairs'\n",
        "    },\n",
        "    {\n",
        "        'color': [169, 164, 131],\n",
        "        'isthing': 0,\n",
        "        'id': 166,\n",
        "        'name': 'tent'\n",
        "    },\n",
        "    {\n",
        "        'color': [225, 199, 255],\n",
        "        'isthing': 0,\n",
        "        'id': 168,\n",
        "        'name': 'towel'\n",
        "    },\n",
        "    {\n",
        "        'color': [137, 54, 74],\n",
        "        'isthing': 0,\n",
        "        'id': 171,\n",
        "        'name': 'wall-brick'\n",
        "    },\n",
        "    {\n",
        "        'color': [135, 158, 223],\n",
        "        'isthing': 0,\n",
        "        'id': 175,\n",
        "        'name': 'wall-stone'\n",
        "    },\n",
        "    {\n",
        "        'color': [7, 246, 231],\n",
        "        'isthing': 0,\n",
        "        'id': 176,\n",
        "        'name': 'wall-tile'\n",
        "    },\n",
        "    {\n",
        "        'color': [107, 255, 200],\n",
        "        'isthing': 0,\n",
        "        'id': 177,\n",
        "        'name': 'wall-wood'\n",
        "    },\n",
        "    {\n",
        "        'color': [58, 41, 149],\n",
        "        'isthing': 0,\n",
        "        'id': 178,\n",
        "        'name': 'water-other'\n",
        "    },\n",
        "    {\n",
        "        'color': [183, 121, 142],\n",
        "        'isthing': 0,\n",
        "        'id': 180,\n",
        "        'name': 'window-blind'\n",
        "    },\n",
        "    {\n",
        "        'color': [255, 73, 97],\n",
        "        'isthing': 0,\n",
        "        'id': 181,\n",
        "        'name': 'window-other'\n",
        "    },\n",
        "    {\n",
        "        'color': [107, 142, 35],\n",
        "        'isthing': 0,\n",
        "        'id': 184,\n",
        "        'name': 'tree-merged'\n",
        "    },\n",
        "    {\n",
        "        'color': [190, 153, 153],\n",
        "        'isthing': 0,\n",
        "        'id': 185,\n",
        "        'name': 'fence-merged'\n",
        "    },\n",
        "    {\n",
        "        'color': [146, 139, 141],\n",
        "        'isthing': 0,\n",
        "        'id': 186,\n",
        "        'name': 'ceiling-merged'\n",
        "    },\n",
        "    {\n",
        "        'color': [70, 130, 180],\n",
        "        'isthing': 0,\n",
        "        'id': 187,\n",
        "        'name': 'sky-other-merged'\n",
        "    },\n",
        "    {\n",
        "        'color': [134, 199, 156],\n",
        "        'isthing': 0,\n",
        "        'id': 188,\n",
        "        'name': 'cabinet-merged'\n",
        "    },\n",
        "    {\n",
        "        'color': [209, 226, 140],\n",
        "        'isthing': 0,\n",
        "        'id': 189,\n",
        "        'name': 'table-merged'\n",
        "    },\n",
        "    {\n",
        "        'color': [96, 36, 108],\n",
        "        'isthing': 0,\n",
        "        'id': 190,\n",
        "        'name': 'floor-other-merged'\n",
        "    },\n",
        "    {\n",
        "        'color': [96, 96, 96],\n",
        "        'isthing': 0,\n",
        "        'id': 191,\n",
        "        'name': 'pavement-merged'\n",
        "    },\n",
        "    {\n",
        "        'color': [64, 170, 64],\n",
        "        'isthing': 0,\n",
        "        'id': 192,\n",
        "        'name': 'mountain-merged'\n",
        "    },\n",
        "    {\n",
        "        'color': [152, 251, 152],\n",
        "        'isthing': 0,\n",
        "        'id': 193,\n",
        "        'name': 'grass-merged'\n",
        "    },\n",
        "    {\n",
        "        'color': [208, 229, 228],\n",
        "        'isthing': 0,\n",
        "        'id': 194,\n",
        "        'name': 'dirt-merged'\n",
        "    },\n",
        "    {\n",
        "        'color': [206, 186, 171],\n",
        "        'isthing': 0,\n",
        "        'id': 195,\n",
        "        'name': 'paper-merged'\n",
        "    },\n",
        "    {\n",
        "        'color': [152, 161, 64],\n",
        "        'isthing': 0,\n",
        "        'id': 196,\n",
        "        'name': 'food-other-merged'\n",
        "    },\n",
        "    {\n",
        "        'color': [116, 112, 0],\n",
        "        'isthing': 0,\n",
        "        'id': 197,\n",
        "        'name': 'building-other-merged'\n",
        "    },\n",
        "    {\n",
        "        'color': [0, 114, 143],\n",
        "        'isthing': 0,\n",
        "        'id': 198,\n",
        "        'name': 'rock-merged'\n",
        "    },\n",
        "    {\n",
        "        'color': [102, 102, 156],\n",
        "        'isthing': 0,\n",
        "        'id': 199,\n",
        "        'name': 'wall-other-merged'\n",
        "    },\n",
        "    {\n",
        "        'color': [250, 141, 255],\n",
        "        'isthing': 0,\n",
        "        'id': 200,\n",
        "        'name': 'rug-merged'\n",
        "    },\n",
        "]\n",
        "\n",
        "# We map the semantic id from 1-200 to contiguous 1-133.\n",
        "for i in range(len(COCO_META)):\n",
        "  COCO_META[i]['id'] = i + 1"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Avk0g2-wo2AO"
      },
      "outputs": [],
      "source": [
        "DatasetInfo = collections.namedtuple(\n",
        "    'DatasetInfo',\n",
        "    'num_classes, label_divisor, thing_list, colormap, class_names')\n",
        "\n",
        "\n",
        "def _coco_label_colormap():\n",
        "  \"\"\"Creates a label colormap used in COCO segmentation benchmark.\n",
        "\n",
        "  See more about COCO dataset at https://cocodataset.org/\n",
        "  Tsung-Yi Lin, et al. \"Microsoft COCO: Common Objects in Context.\" ECCV. 2014.\n",
        "\n",
        "  Returns:\n",
        "    A 2-D numpy array with each row being mapped RGB color (in uint8 range).\n",
        "  \"\"\"\n",
        "  colormap = np.zeros((256, 3), dtype=np.uint8)\n",
        "  for category in COCO_META:\n",
        "    colormap[category['id']] = category['color']\n",
        "  return colormap\n",
        "\n",
        "\n",
        "def _coco_class_names():\n",
        "  return ('void',) + tuple([x['name'] for x in COCO_META])\n",
        "\n",
        "\n",
        "def coco_dataset_information():\n",
        "  return DatasetInfo(\n",
        "      num_classes=134,\n",
        "      label_divisor=256,\n",
        "      thing_list=tuple(range(1, 81)),\n",
        "      colormap=_coco_label_colormap(),\n",
        "      class_names=_coco_class_names())\n",
        "\n",
        "\n",
        "def perturb_color(color, noise, used_colors, max_trials=50, random_state=None):\n",
        "  \"\"\"Pertrubs the color with some noise.\n",
        "\n",
        "  If `used_colors` is not None, we will return the color that has\n",
        "  not appeared before in it.\n",
        "\n",
        "  Args:\n",
        "    color: A numpy array with three elements [R, G, B].\n",
        "    noise: Integer, specifying the amount of perturbing noise (in uint8 range).\n",
        "    used_colors: A set, used to keep track of used colors.\n",
        "    max_trials: An integer, maximum trials to generate random color.\n",
        "    random_state: An optional np.random.RandomState. If passed, will be used to\n",
        "      generate random numbers.\n",
        "\n",
        "  Returns:\n",
        "    A perturbed color that has not appeared in used_colors.\n",
        "  \"\"\"\n",
        "  if random_state is None:\n",
        "    random_state = np.random\n",
        "\n",
        "  for _ in range(max_trials):\n",
        "    random_color = color + random_state.randint(\n",
        "        low=-noise, high=noise + 1, size=3)\n",
        "    random_color = np.clip(random_color, 0, 255)\n",
        "\n",
        "    if tuple(random_color) not in used_colors:\n",
        "      used_colors.add(tuple(random_color))\n",
        "      return random_color\n",
        "\n",
        "  print('Max trial reached and duplicate color will be used. Please consider '\n",
        "        'increase noise in `perturb_color()`.')\n",
        "  return random_color\n",
        "\n",
        "\n",
        "def color_panoptic_map(panoptic_prediction, dataset_info, perturb_noise):\n",
        "  \"\"\"Helper method to colorize output panoptic map.\n",
        "\n",
        "  Args:\n",
        "    panoptic_prediction: A 2D numpy array, panoptic prediction from deeplab\n",
        "      model.\n",
        "    dataset_info: A DatasetInfo object, dataset associated to the model.\n",
        "    perturb_noise: Integer, the amount of noise (in uint8 range) added to each\n",
        "      instance of the same semantic class.\n",
        "\n",
        "  Returns:\n",
        "    colored_panoptic_map: A 3D numpy array with last dimension of 3, colored\n",
        "      panoptic prediction map.\n",
        "    used_colors: A dictionary mapping semantic_ids to a set of colors used\n",
        "      in `colored_panoptic_map`.\n",
        "  \"\"\"\n",
        "  if panoptic_prediction.ndim != 2:\n",
        "    raise ValueError('Expect 2-D panoptic prediction. Got {}'.format(\n",
        "        panoptic_prediction.shape))\n",
        "\n",
        "  semantic_map = panoptic_prediction // dataset_info.label_divisor\n",
        "  instance_map = panoptic_prediction % dataset_info.label_divisor\n",
        "  height, width = panoptic_prediction.shape\n",
        "  colored_panoptic_map = np.zeros((height, width, 3), dtype=np.uint8)\n",
        "\n",
        "  used_colors = collections.defaultdict(set)\n",
        "  # Use a fixed seed to reproduce the same visualization.\n",
        "  random_state = np.random.RandomState(0)\n",
        "\n",
        "  unique_semantic_ids = np.unique(semantic_map)\n",
        "  for semantic_id in unique_semantic_ids:\n",
        "    semantic_mask = semantic_map == semantic_id\n",
        "    if semantic_id in dataset_info.thing_list:\n",
        "      # For `thing` class, we will add a small amount of random noise to its\n",
        "      # correspondingly predefined semantic segmentation colormap.\n",
        "      unique_instance_ids = np.unique(instance_map[semantic_mask])\n",
        "      for instance_id in unique_instance_ids:\n",
        "        instance_mask = np.logical_and(semantic_mask,\n",
        "                                       instance_map == instance_id)\n",
        "        random_color = perturb_color(\n",
        "            dataset_info.colormap[semantic_id],\n",
        "            perturb_noise,\n",
        "            used_colors[semantic_id],\n",
        "            random_state=random_state)\n",
        "        colored_panoptic_map[instance_mask] = random_color\n",
        "    else:\n",
        "      # For `stuff` class, we use the defined semantic color.\n",
        "      colored_panoptic_map[semantic_mask] = dataset_info.colormap[semantic_id]\n",
        "      used_colors[semantic_id].add(tuple(dataset_info.colormap[semantic_id]))\n",
        "  return colored_panoptic_map, used_colors\n",
        "\n",
        "\n",
        "def vis_segmentation(image,\n",
        "                     panoptic_prediction,\n",
        "                     dataset_info,\n",
        "                     perturb_noise=60):\n",
        "  \"\"\"Visualizes input image, segmentation map and overlay view.\"\"\"\n",
        "  plt.figure(figsize=(30, 20))\n",
        "  grid_spec = gridspec.GridSpec(2, 2)\n",
        "\n",
        "  ax = plt.subplot(grid_spec[0])\n",
        "  plt.imshow(image)\n",
        "  plt.axis('off')\n",
        "  ax.set_title('input image', fontsize=20)\n",
        "\n",
        "  ax = plt.subplot(grid_spec[1])\n",
        "  panoptic_map, used_colors = color_panoptic_map(panoptic_prediction,\n",
        "                                                 dataset_info, perturb_noise)\n",
        "  plt.imshow(panoptic_map)\n",
        "  plt.axis('off')\n",
        "  ax.set_title('panoptic map', fontsize=20)\n",
        "\n",
        "  ax = plt.subplot(grid_spec[2])\n",
        "  plt.imshow(image)\n",
        "  plt.imshow(panoptic_map, alpha=0.7)\n",
        "  plt.axis('off')\n",
        "  ax.set_title('panoptic overlay', fontsize=20)\n",
        "\n",
        "  ax = plt.subplot(grid_spec[3])\n",
        "  max_num_instances = max(len(color) for color in used_colors.values())\n",
        "  # RGBA image as legend.\n",
        "  legend = np.zeros((len(used_colors), max_num_instances, 4), dtype=np.uint8)\n",
        "  class_names = []\n",
        "  for i, semantic_id in enumerate(sorted(used_colors)):\n",
        "    legend[i, :len(used_colors[semantic_id]), :3] = np.array(\n",
        "        list(used_colors[semantic_id]))\n",
        "    legend[i, :len(used_colors[semantic_id]), 3] = 255\n",
        "    if semantic_id \u003c dataset_info.num_classes:\n",
        "      class_names.append(dataset_info.class_names[semantic_id])\n",
        "    else:\n",
        "      class_names.append('ignore')\n",
        "\n",
        "  plt.imshow(legend, interpolation='nearest')\n",
        "  ax.yaxis.tick_left()\n",
        "  plt.yticks(range(len(legend)), class_names, fontsize=15)\n",
        "  plt.xticks([], [])\n",
        "  ax.tick_params(width=0.0, grid_linewidth=0.0)\n",
        "  plt.grid('off')\n",
        "  plt.show()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1ly6p6M2o8SF"
      },
      "source": [
        "### Select a pretrained model"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "peo7LUTtulpQ"
      },
      "outputs": [],
      "source": [
        "MODEL_NAME = 'resnet50_kmax_deeplab_coco_train'  # @param ['resnet50_kmax_deeplab_coco_train','axial_resnet50_kmax_deeplab_coco_train','convnext_tiny_kmax_deeplab_coco_train','convnext_small_kmax_deeplab_coco_train','convnext_base_kmax_deeplab_coco_train','convnext_large_kmax_deeplab_coco_train','convnext_large_kmax_deeplab_coco_train_unlabeled']\n",
        "\n",
        "\n",
        "_MODELS = ('resnet50_kmax_deeplab_coco_train',\n",
        "           'axial_resnet50_kmax_deeplab_coco_train',\n",
        "           'convnext_tiny_kmax_deeplab_coco_train',\n",
        "           'convnext_small_kmax_deeplab_coco_train',\n",
        "           'convnext_base_kmax_deeplab_coco_train',\n",
        "           'convnext_large_kmax_deeplab_coco_train',\n",
        "           'convnext_large_kmax_deeplab_coco_train_unlabeled'\n",
        "           )\n",
        "_DOWNLOAD_URL_PATTERN = 'https://storage.googleapis.com/gresearch/tf-deeplab/saved_model/%s.tar.gz'\n",
        "\n",
        "_MODEL_NAME_TO_URL_AND_DATASET = {\n",
        "    model: (_DOWNLOAD_URL_PATTERN % model, coco_dataset_information())\n",
        "    for model in _MODELS\n",
        "}\n",
        "\n",
        "MODEL_URL, DATASET_INFO = _MODEL_NAME_TO_URL_AND_DATASET[MODEL_NAME]\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "UjYwP1Sjo4dd"
      },
      "outputs": [],
      "source": [
        "model_dir = tempfile.mkdtemp()\n",
        "\n",
        "download_path = os.path.join(model_dir, MODEL_NAME + '.gz')\n",
        "urllib.request.urlretrieve(MODEL_URL, download_path)\n",
        "\n",
        "!tar -xzvf {download_path} -C {model_dir}\n",
        "\n",
        "LOADED_MODEL = tf.saved_model.load(os.path.join(model_dir, MODEL_NAME))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "umpwnn4etG6z"
      },
      "source": [
        "### Run on sample images"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "6552FXlAOHnX"
      },
      "outputs": [],
      "source": [
        "# Required, upload an image from your local machine.\n",
        "\n",
        "uploaded = files.upload()\n",
        "\n",
        "if not uploaded:\n",
        "  raise AssertionError('Please upload one image')\n",
        "elif len(uploaded) == 1:\n",
        "  UPLOADED_FILE = list(uploaded.keys())[0]\n",
        "else:\n",
        "  raise AssertionError('Please upload one image at a time')"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "bsQ7Oj7jtHDz"
      },
      "outputs": [],
      "source": [
        "with tf.io.gfile.GFile(UPLOADED_FILE, 'rb') as f:\n",
        "  im = np.array(Image.open(f))\n",
        "\n",
        "output = LOADED_MODEL(tf.cast(im, tf.uint8))\n",
        "vis_segmentation(im, output['panoptic_pred'][0], DATASET_INFO)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "k2-EDpBCwmpq"
      },
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "colab": {
      "collapsed_sections": [],
      "name": "DeepLab_COCO_Demo.ipynb",
      "private_outputs": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
